# coding: utf-8
import pysequoiadb
from pysequoiadb import client
from pysequoiadb.error import SDBBaseError
from pysequoiadb.error import SDBEndOfCursor

import time
import base64
from io import StringIO
import json

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

import redis
from keras.models import load_model
import argparse
from scipy import misc
import numpy as np
import os
import cv2
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

flags = argparse.ArgumentParser()
flags.add_argument('--gpu', default=0, type=float, help='gpu number')
flags.add_argument('--redis_ip', help='Redis IP.')
flags.add_argument('--redis_port', type=int, help='Redis IP port.')
flags.add_argument('--redis_db', default=0, type=int, help='Redis db number.')
flags.add_argument('--sdb_ip', help='The ip of SDB.')
flags.add_argument('--sdb_port', help='The port of SDB.')
flags.add_argument('--channel', help='camera channel.')
flags.add_argument('--sdb_status_col', default='station.rstp', help='Select status from sdb collection.')
flags.add_argument('--redis_list_length', default=100, type=int, help='')
flags.add_argument('--station', help='model select')
flags.add_argument('--theshold', default=0.5, type=float, help='The lower bound tolerance down  of distince.')
flags.add_argument('--redis_key', help='The key of redis.')
FLAGS = flags.parse_args()

# 使用GPU
gpun = FLAGS.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpun)

def sdb_connect(ip, port):
	# connect to local db, using default args value.
	# host= 'localhost', port= 11810, user= '', password= ''
	try:
		db = client(ip,port,"peter","Foxconn99.")
		return db
	except SDBBaseError as e:
		pysequoiadb._print(e)
		del db
		exit()

def get_sdb_status(col_name, db, db_num, channel):
	select_sql = "select status from {} where db={} and chn={}".format(col_name, db_num, channel)
	#print select_sql
	cursor = db.exec_sql(select_sql)
	status = 1
	try:
		while True:
			try:
				bson = cursor.next()
				status = bson.get("status")
			except SDBEndOfCursor:
				break
	except:
		print("Warning: SDB is Down.")
	finally:
		cursor.close()
	return status

def do_crop(img):
	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
	th2 = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 2)
	contours, hierarchy = cv2.findContours(th2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
	i = 0
	maxarea = -1
	roci = (0, 0, 0, 0)
	for contour in contours:
		x, y, w, h = cv2.boundingRect(contour)
		if (not abs(h - y) < 5) and not (abs(h - y) > 50):  # 比例>0.2
			# cv2.rectangle(self.img, (x, y), (x + w, y + h), (0, 255, 0), 2)
			tarea = w * h
			if maxarea < tarea:
				maxarea = tarea
				roci = (x, y, w, h)
	# print("矩形宽，高", roci[2], '    ', roci[3])
	try:
		img__ = img[roci[1]:roci[1] + roci[3], roci[0]:roci[0] + roci[2]]
		img__ = cv2.resize(img__,dsize=(123,126))
	except:
		print("size is error")
		img__ = None
		#img__ = cv2.resize(gray,dsize=(123,126))
	return img__


def resize_image(images, ROWS, COLS, CHANNELS):
	data = np.ndarray((len(images), ROWS, COLS, CHANNELS), dtype=np.uint8)
	for i, image in enumerate(images):
		image = do_crop(image)
		if image is not None:
			data[i] = misc.imresize(image, (ROWS, COLS), interp='bilinear')
	return data

if __name__ == '__main__':
	program_starttime = time.time()
	# Link to redis
	ip = FLAGS.redis_ip
	redis_port = FLAGS.redis_port
	db_num = FLAGS.redis_db
	password = "Foxconn88"
	conn = redis.StrictRedis(host=ip, port=redis_port, password=password, db=db_num)

	#Link to sdb
	sdb_ip = FLAGS.sdb_ip
	sdb_port = FLAGS.sdb_port
	db = sdb_connect(sdb_ip, sdb_port)
	sdb_status_col = FLAGS.sdb_status_col

	channel = FLAGS.channel
	redis_list_length = FLAGS.redis_list_length

	#tolerancedown = FLAGS.theshold_2
	#toleranceup = FLAGS.theshold_3
	tolerance=FLAGS.theshold
	redis_key = FLAGS.redis_key

	#限定GPU最小使用
	config = tf.ConfigProto()
	config.gpu_options.allow_growth=True
	session = tf.Session(config=config)
	KTF.set_session(session)

	# 加载模型
	path, _ = os.path.split(os.path.realpath(__file__))
	model_index = FLAGS.station
#    model_path = "/hadoop/station_photo_determine/model/newmodel/resnet_34_1020.h5"
	model_path = path + '/model/' + model_index + '_model.h5'
	model = load_model(model_path)
	print('model load ok')
	try:
		while True:
			starttime = time.time()
			# 检查终止条件
			status = get_sdb_status(sdb_status_col, db, db_num, channel)
			if status == 0:
				db.disconnect()
				del db
				break

			# 捕获redis异常
			try:
				len_redis_key = conn.llen(redis_key)
			except:
				print("redis连接问题，请检查redis状态和连接IP.")
				#cl = db.get_collection(sdb_status_col)
				#condition = {'chn': int(channel), 'db': db_num}
				#rule = {"$set": {"status": 0}}
				#cl.update(rule=rule, condition=condition)
				#db.disconnect()
				#del db
				continue

			# 读取redis数据
			images_byte_str = []
			image_arr = []
			labels = []
			flags = []
			uploadtimes = []
			live_status = 0
			died_status = 1
			try:
				# 读取数据，更新状态为1，暂时不能写入数据到redis，读取完成在更新状态为0，可写数据
				conn.set('wsread_status_{}'.format(channel), died_status)
				record_dict = [eval(record) for record in conn.lrange(redis_key, 0, redis_list_length-1)]
				#conn.delete(redis_key)
				conn.ltrim(redis_key,0,-1-(len(record_dict)))
				conn.set('wsread_status_{}'.format(channel), live_status)

				for record in record_dict:
					images_byte_str.append(record["imgstr"])
					uploadtimes.append(record["time"])
					flags.append(record["flag"])
					labels.append(record["label"])
				images_byte_str = np.array(images_byte_str)
				uploadtimes = np.array(uploadtimes)
				flags = np.array(flags)
				labels = np.array(labels)

				exclude_index = np.ones_like(images_byte_str, dtype=np.bool)

				for i, image_byte in enumerate(images_byte_str):
					try:
						decode_byte = base64.b64decode(image_byte)
						image = misc.imread(StringIO(decode_byte))[:,:,:3]
						image_arr.append(image)
					except:
						exclude_index[i] = False
						continue

				uploadtimes = uploadtimes[exclude_index]
				labels = labels[exclude_index]
				flags = flags[exclude_index]

			except:
				print("Have some errors when read pictures from the redis.")
				continue

			if len(image_arr) == 0:
				# print("No pictures have been selected from the redis by flag.")
				continue

			# 统一图片大小
			ROWS = 297
			COLS = 396
			CHANNELS = 3
			resize_img =  resize_image(image_arr,ROWS,COLS,CHANNELS)

			# 预测设备好坏
			model_start_time = time.time()
			predictions = model.predict(resize_img, verbose=0)
			model_end_time = time.time()
			print("Model prediction use time: {}".format(model_end_time-model_start_time))

			print(predictions)
			#获取最大预测值的索引
			predictionsIndex=np.argmax(predictions,axis=1)
			maxPredictions=[]
			cate = ['ONE','TWO','THREE','FOUR','FIVE','OTHER']
			results=[]
			#最大预测概率大于阈值确定分类,否则忽略
			for i in range(len(predictions)):
				predic = predictions[i][predictionsIndex[i]]
				if predic < tolerance:
					results.append('NONE')
				else:
					results.append(cate[predictionsIndex[i]])
		#最大预测值
				maxPredictions.append(predic)

			print(predictionsIndex)
			print(maxPredictions)
			print(results)

			redis_dict = {}
			redis_dict["channel"] = channel
			redis_dict["prediction"] =np.array( maxPredictions).tolist() #predictions.tolist()
			redis_dict["time"] = uploadtimes.tolist()
			redis_dict["label"] = labels.tolist()
			redis_dict["uuid"] = flags.tolist()
			redis_dict["status"] = np.array(results).tolist() #results.tolist()
			json_redis_dict = json.dumps(redis_dict)

			key = "ws_json_redis_dict_{}".format(channel)
			conn.lpush(key, str(json_redis_dict))
			len_key = conn.llen(key)
			if len_key > 100:
				conn.ltrim(key, 0, 99)
			# conn.set('img_id_{}'.format(channel), img_id)
			print("runnig time: %f" % (time.time() - starttime))

	except KeyboardInterrupt:
		# Need to release client whether it connected db server successfully or not
		print("Programe end.")
		db.disconnect()
		del db
		endtime = time.time()
		print('Running time: %s Seconds' % (endtime - program_starttime))


